from hybrid_search import hybrid_search
from rag_evaluation import evaluate_retrieval_relevance


def find_best_weights(test_questions):
    """
    自动寻找最优的混合检索权重

    参数:
        test_questions: 测试问题列表

    返回:
        最佳的es_weight和milvus_weight
    """
    best_score = 0
    best_weights = (0.3, 0.7)  # 默认权重

    # 尝试不同的权重组合（es_weight从0.1到0.5）
    for es_w in [0.1, 0.2, 0.3, 0.4, 0.5]:
        milvus_w = 1 - es_w  # 权重和为1
        print(f"测试权重：ES={es_w}, Milvus={milvus_w}")

        # 收集该权重下的检索结果
        search_results = []
        for question in test_questions:
            results = hybrid_search(question, es_weight=es_w, milvus_weight=milvus_w)
            search_results.append(results)

        # 评估检索相关性
        score = evaluate_retrieval_relevance(test_questions, search_results)

        # 记录最佳权重
        if score > best_score:
            best_score = score
            best_weights = (es_w, milvus_w)

    print(f"\n最佳权重：ES={best_weights[0]}, Milvus={best_weights[1]}, 分数={best_score:.2f}")
    return best_weights


def optimize_chunk_parameters(test_questions):
    """
    简单优化文本切分参数（示例）
    """
    # 实际应用中可以类似上面的方法，测试不同的chunk_size和chunk_overlap
    print("\n建议切分参数：")
    print("1. 保险条款较长，建议chunk_size=300")
    print("2. 为保持条款完整，建议chunk_overlap=80")
    return (300, 80)


# 测试优化代码
if __name__ == "__main__":
    # 测试问题
    test_questions = [
        "平安福2024版的保障范围是什么？",
        "平安福2024版的投保年龄限制是多少？"
    ]

    # 1. 寻找最佳检索权重
    best_es_w, best_milvus_w = find_best_weights(test_questions)

    # 2. 优化文本切分参数
    best_chunk_size, best_chunk_overlap = optimize_chunk_parameters(test_questions)
